import json
import os.path


def get_data(data_type, callback=None):
    assert data_type in ['train', 'dev', 'test']
    dir_name = os.path.dirname(__file__)
    input_file = os.path.join(dir_name, 'data', data_type + '.char.bmes')
    features = []
    with open(input_file, 'r', encoding="utf-8") as f:
        words = []
        labels = []
        tmp_line = None
        for line in f:
            if line.startswith("-DOCSTART-") or line == "" or line == "\n":
                if words:
                    tmp_line = {"words": words, "label": labels}
                    words = []
                    labels = []
            else:
                splits = line.split(" ")
                words.append(splits[0])
                if len(splits) > 1:
                    labels.append(splits[-1].replace("\n", ""))
                else:
                    # Examples could have no label for mode = "test"
                    labels.append("O")
            if tmp_line:
                text = "".join(tmp_line['words'])
                label_list = tmp_line['label']
                line = {"text": text, "labels": label_list}
                if callback is not None:
                    line = callback(line=line)
                features.append(line)
                tmp_line = None
    return features
if __name__ == '__main__':
    data = get_data('train')
    print(len(data))
